#!/usr/bin/env python

import os
import pexpect
import sys
import re
import shutil

# Put this script in the firmadyne path downloadable from
# https://github.com/firmadyne/firmadyne

#Configurations - change this according to your system
#firmadyne_path = "/home/ubuntu/iot/firmware-analysis-toolkit/firmadyne"
#firmadyne_path="/home/ubuntu/iot/firmware-analysis-toolkit/firmadyne/"
firmadyne_path="/home/h4lo/iot/env/firmadyne/"
binwalk_path = "/usr/local/bin/binwalk"
root_pass = "password"
firmadyne_pass = "firmadyne"
user_name = "h4lo"
local_tap_ip = "10.0.0.2"
ifconfig_cmd = "ifconfig eth0 10.0.0.1/24"

def show_banner():
    print """
                               __           _
                              / _|         | |
                             | |_    __ _  | |_
                             |  _|  / _` | | __|
                             | |   | (_| | | |_
                             |_|    \__,_|  \__|

                Welcome to the Firmware Analysis Toolkit - v0.2
    Offensive IoT Exploitation Training  - http://offensiveiotexploitation.com
                  By Attify - https://attify.com  | @attifyme
    """


def get_info():
    if len(sys.argv) == 2:
        firm_name = sys.argv[1]
        print "[?] Enter the name or absolute path of the firmware you want to analyse : " + firm_name
    else:
        firm_name = raw_input("[?] Enter the name or absolute path of the firmware you want to analyse : ")
    firm_brand = raw_input("[?] Enter the brand of the firmware : ")
    return (firm_name, firm_brand)


def run_extractor(firm_name, firm_brand):
    print "[+] Now going to extract the firmware. Hold on.."
    print "[+] Firmware : " + firm_name
    print "[+] Brand : " + firm_brand
    extractor_cmd = firmadyne_path + "/sources/extractor/extractor.py -b " + firm_brand + " -sql 127.0.0.1 -np -nk " + "\""+ firm_name + "\"" + " images "
    child = pexpect.spawn(extractor_cmd, timeout=None)
    child.expect("Database Image ID: ")
    image_id = child.readline().strip()
    print "[+] Database image ID : " + image_id
    child.expect(pexpect.EOF)
    return image_id


def identify_arch(image_id):
    print "[+] Identifying architecture"
    identfy_arch_cmd = firmadyne_path + "/scripts/getArch.sh ./images/" + image_id + ".tar.gz"
    child = pexpect.spawn(identfy_arch_cmd)
    child.expect(":")
    arch = child.readline().strip()
    print "[+] Architecture : " + arch
    child.expect("Password for user firmadyne: ")
    child.sendline(firmadyne_pass)
    child.expect(pexpect.EOF)
    return arch


def tar2db(image_id):
    print "[+] Storing filesystem in database"
    tar2db_cmd = firmadyne_path + "/scripts/tar2db.py -i " + image_id + " -f " + firmadyne_path + "/images/" + image_id + ".tar.gz"
    output_tar2db = pexpect.run(tar2db_cmd)

    if "already exists" in output_tar2db:
        print "[!] Filesystem already exists"


def make_image(arch, image_id):
    print "[+] Building QEMU disk image"
    makeimage_cmd = "sudo " + firmadyne_path + "/scripts/makeImage.sh " + image_id + " " + arch
    child = pexpect.spawn(makeimage_cmd)
    child.sendline(root_pass)
    child.expect(pexpect.EOF)


def setup_network(arch, image_id):
    print "[+] Setting up the network connection, please standby"
    network_cmd = "sudo " + firmadyne_path + "/scripts/inferNetwork.sh " + image_id + " " + arch
    child = pexpect.spawn(network_cmd)
    child.sendline(root_pass)
    child.expect("Interfaces:", timeout=None)
    interfaces = child.readline().strip()
    print "[+] Network interfaces : " + interfaces
    child.expect(pexpect.EOF)
    return interfaces

def final_run(image_id):
    print "[+] Running the firmware finally"
    run_cmd = "sudo " + firmadyne_path + "/scratch/" + image_id + "/run.sh"
    print "[+] command line : " + run_cmd
    raw_input("[*] Press ENTER to run the firmware...")
    child = pexpect.spawn(run_cmd)
    child.sendline(root_pass)
    child.interact()


def change_config_permiss(file_name):
    run_cmd = "sudo chmod 777 " + str(file_name)
    child = pexpect.spawn(run_cmd)
    child.sendline(root_pass)
    res = child.readline()

def change_network_to_tap(image_id):
    run_file = firmadyne_path + "/scratch/" + image_id + "/run.sh"
    #print "file: " + str(run_file)
    tap_config = "-net nic -net tap,ifname=tap0,script=no,downscript=no"
    change_config_permiss(run_file)
    orig_content = ""
    try:
	with open(run_file,mode="r") as fr:
	    orig_content = fr.read()

	    fr.close()
	#print orig_content

	with open(run_file,mode="w") as fw:
	    orig_net = re.search("-nographic (.[\s\S*]*)",orig_content).group(1)
	    #print orig_net
	    fw.write(orig_content.replace(orig_net,tap_config))

	    fw.close()

	return True
    except Exception as e:
	print str(e)
	return False

    return True

def turn_to_tap():
    try:
	run_cmd = "sudo tunctl -t tap0 -u "+str(user_name)

	child = pexpect.spawn(run_cmd)
	child.sendline(root_pass)
	run_cmd = "sudo ifconfig tap0 {ip}/8".format(ip=local_tap_ip)
	child = pexpect.spawn(run_cmd)
	child.sendline(root_pass)
   	return True
    except:
	return False

    return True


def add_ipconfig(image_id):
    run_cmd = 'sed -i "s/#\!\/bin\/sh/#\!\/bin\/sh\n{CONFIG_IP}\ntelnetd -l \/bin\/sh -p 2309/g"'.format(CONFIG_IP=ifconfig_cmd)

    if(os.path.exists("./scratch/{id}/image/etc/init.d/rcS".format(id=image_id))):
	pexpect.spawn(run_cmd + " " +str("./scratch/{id}/image/etc/init.d/rcS".format(id=image_id)))
    elif(os.path.exists("./scratch/{id}/image/etc/rcS".format(id=image_id))):
	pexpect.spawn(run_cmd + " " +str("./scratch/{id}/image/etc/rcS".format(id=image_id)))
    elif(os.path.exists("./scratch/{id}/image/etc/preinit".format(id=image_id))):
	pexpect.spawn(run_cmd + " " +str("./scratch/{id}/image/etc/preinit".format(id=image_id)))
    else:
	return False

    return True



def main():
    show_banner()
    firm_name, firm_brand = get_info()
    image_id = run_extractor(firm_name, firm_brand)

    if image_id == "":
        print "[!] Something went wrong"
    else:
        arch = identify_arch(image_id)
    #    tar2db(image_id)
        make_image(arch, image_id)
        interface = setup_network(arch, image_id)
    #    final_run(image_id)

	if len(str(interface[0]))<5:
            if(change_network_to_tap(image_id)):
		print "[+] Modify `run.sh` file success!"
	    else:
		print "[+] Modify `run.sh` file fail!"
		exit(1)
	    try:
		turn_to_tap()
	    except:
		print "[-] Turn to tap pattern fail!!"

		pass
	    print "[+] Router IP: 10.0.0.1"

	    if(not add_ipconfig(image_id)):
		print("[+] Path or init file not found!")
	    else:
		pass

	else:
		print "[+] Router IP: " + str(interface)
	print "[+] Done!Now to run the command :`sudo ./scratch/{iid}/run.sh`".format(iid=str(image_id))

def test_main():
    change_network_to_tap('1')
    turn_to_tap()

if __name__ == "__main__":
    main()
    #test_main()